import gymnasium as gym
from gymnasium import spaces
import numpy as np
import mujoco
import glfw

class YmbotEPIDEnv(gym.Env):
    def __init__(self, render_mode=None, window_width=1200, window_height=900):
        super(YmbotEPIDEnv, self).__init__()
        # 加载 ymbot_e 的 Mujoco 模型
        self.model = mujoco.MjModel.from_xml_path('/home/sh/catkin_ws/src/ymbot_e_description/mjcf/ymbot_e_description.xml')
        self.data = mujoco.MjData(self.model)
        self.model.opt.timestep = 0.001  # 设置底层仿真时间为 0.001 秒
        self.render_mode = render_mode
        # 定义自定义环境的动作空间，只训练最后一个关节的力矩
        self.action_space = spaces.Box(low=-200, high=200, shape=(1,), dtype=np.float32)
        # 定义自定义环境的观测空间，30 个关节的角度，每个角度范围为 [-10, 10]
        self.observation_space = spaces.Box(low=-3.14, high=3.14, shape=(30,), dtype=np.float32)
        # 目标值曲线，默认函数
        self.setpoint_curve = self.default_target_curve
        # 当前状态
        self.state = np.zeros(30)
        # 上一次的误差
        self.last_error = np.zeros(30)
        # 积分项
        self.integral = np.zeros(30)
        # 时间步
        self.timestep = 0
        self.control_output = np.zeros(30)  # 初始化控制输出
        self.mouse_button_left = False
        self.mouse_button_middle = False
        self.mouse_button_right = False
        self.last_mouse_x = 0
        self.last_mouse_y = 0
        self.scene = None
        self.camera = mujoco.MjvCamera()
        self.window_width = window_width
        self.window_height = window_height
        # 初始化误差记录
        self.error_history = []

    def default_target_curve(self, t):
        """默认目标曲线函数"""
        curve = np.zeros(30)
        curve[28] = 20/180*np.pi * np.sin((t-1) * 0.01)  
        curve[29] = 20/180*np.pi * np.sin((t-1) * 0.01)  
        return curve

    def set_target_curve(self, curve_function):
        """设置目标曲线"""
        self.setpoint_curve = curve_function

    def reset(self, seed=None, options=None):
        # 重置 ymbot_e 环境
        mujoco.mj_resetData(self.model, self.data)
        # 随机初始化当前状态
        self.state = np.random.uniform(-1, 1, size=30)
        # 重置误差和积分项
        self.last_error = np.zeros(30)
        self.integral = np.zeros(30)
        # 重置时间步
        self.timestep = 0
        return self.state, {}

    def step(self, action):
        # 注释掉 PID 参数相关代码
        # kp = np.full(30, 200)
        # kp[29] = action[0]
        # ki = np.zeros(30)
        # kd = np.zeros(30)
        
        # 更新目标值
        self.setpoint = self.setpoint_curve(self.timestep)
        
        # 计算误差、积分和导数项
        error = self.setpoint - self.state
        # self.integral += error
        # derivative = error - self.last_error
        self.last_error = error
        # 计算控制量
        # self.control_output = kp * error + ki * self.integral + kd * derivative
        
        # 直接使用动作作为控制量，只训练第30个关节的控制器力矩
        self.control_output[29] = action[0]
        
        # 将控制量应用到 ymbot_e 环境的所有关节
        self.data.ctrl[:] = np.clip(self.control_output, -500, 500)
        
        # 更新一步仿真
        mujoco.mj_step(self.model, self.data)
        
        # 更新状态
        self.state = self.data.qpos[:30]  # 使用观测空间中的前30个元素作为关节角度
        # 奖励为负的绝对误差之和
        reward = -np.sum(np.abs(error))
        
        # 记录误差
        self.error_history.append(np.abs(error[29]))
        if len(self.error_history) > 100:  # 记录最近100个误差
            self.error_history.pop(0)
        
        # 使用误差平均值判断任务是否完成
        average_error = np.mean(self.error_history)
        done = (average_error < 0.001) | (average_error > 0.1) # 误差小于 0.001 或大于 0.1 时任务完成
        
        # 增加时间步
        self.timestep += 1
        
        # 奖励为负的绝对误差之和，加上误差减少的奖励
        if len(self.error_history) > 1:
            error_reduction_reward = self.error_history[-2] - self.error_history[-1]
        else:
            error_reduction_reward = 0
        reward = -np.sum(np.abs(error)) + error_reduction_reward
        
        return self.state, reward, done, False, {}

    def get_control_output(self):
        """获取控制输出"""
        return self.control_output

    def get_simulation_timestep(self):
        """获取 Mujoco 仿真时的底层步进时间"""
        return self.model.opt.timestep

    def mouse_button_callback(self, window, button, action, mods):
        if button == glfw.MOUSE_BUTTON_LEFT:
            self.mouse_button_left = (action == glfw.PRESS)
        elif button == glfw.MOUSE_BUTTON_MIDDLE:
            self.mouse_button_middle = (action == glfw.PRESS)
        elif button == glfw.MOUSE_BUTTON_RIGHT:
            self.mouse_button_right = (action == glfw.PRESS)
        self.last_mouse_x, self.last_mouse_y = glfw.get_cursor_pos(window)

    def mouse_move_callback(self, window, xpos, ypos):
        if not self.mouse_button_left and not self.mouse_button_middle and not self.mouse_button_right:
            return

        dx = xpos - self.last_mouse_x
        dy = ypos - self.last_mouse_y
        self.last_mouse_x = xpos
        self.last_mouse_y = ypos

        width, height = glfw.get_window_size(window)
        mod_shift = (glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS or
                     glfw.get_key(window, glfw.KEY_RIGHT_SHIFT) == glfw.PRESS)

        if self.mouse_button_right:
            action = mujoco.mjtMouse.mjMOUSE_MOVE_H if mod_shift else mujoco.mjtMouse.mjMOUSE_MOVE_V
        elif self.mouse_button_left:
            action = mujoco.mjtMouse.mjMOUSE_ROTATE_H if mod_shift else mujoco.mjtMouse.mjMOUSE_ROTATE_V
        else:
            action = mujoco.mjtMouse.mjMOUSE_ZOOM

        mujoco.mjv_moveCamera(self.model, action, dx / height, dy / height, self.scene, self.camera)

    def scroll_callback(self, window, xoffset, yoffset):
        mujoco.mjv_moveCamera(self.model, mujoco.mjtMouse.mjMOUSE_ZOOM, 0, -0.05 * yoffset, self.scene, self.camera)

    def render(self):
        if self.render_mode == 'human':
            if not hasattr(self, 'context'):
                if not glfw.init():
                    raise Exception("GLFW can't be initialized")
                self.window = glfw.create_window(self.window_width, self.window_height, "Ymbot E PID Env", None, None)
                if not self.window:
                    glfw.terminate()
                    raise Exception("GLFW window can't be created")
                glfw.make_context_current(self.window)
                self.context = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
                glfw.set_mouse_button_callback(self.window, self.mouse_button_callback)
                glfw.set_cursor_pos_callback(self.window, self.mouse_move_callback)
                glfw.set_scroll_callback(self.window, self.scroll_callback)
                self.scene = mujoco.MjvScene(self.model, maxgeom=1000)
            viewport = mujoco.MjrRect(0, 0, self.window_width, self.window_height)
            mujoco.mjv_updateScene(self.model, self.data, mujoco.MjvOption(), None, self.camera, mujoco.mjtCatBit.mjCAT_ALL, self.scene)
            mujoco.mjr_render(viewport, self.scene, self.context)
            glfw.swap_buffers(self.window)
            glfw.poll_events()

    def close(self):
        if hasattr(self, 'window'):
            glfw.destroy_window(self.window)
            glfw.terminate()

if __name__ == "__main__":
    env = YmbotEPIDEnv(render_mode='human', window_width=2560, window_height=1440)
    
    def target_curve(t):
        curve = np.zeros(30)
        # curve[28] = 20/180*np.pi * np.sin((t-1) * 0.01)  
        curve[29] = 20/180*np.pi * np.sin((t-1) * 0.01)  
        return curve
    
    env.set_target_curve(target_curve)  # 设置目标曲线
    
    state, _ = env.reset()


    print(f"Simulation Timestep: {env.get_simulation_timestep()}")  # 打印仿真步进时间
    with open('/home/sh/catkin_ws/src/ymbot_e_control/data/ymbot_e_ppo_data.txt', 'w') as file:
        for timestep in range(1000):
            # 修改动作初始化代码，只设置最后一个关节的控制器力矩
            action = np.array([10.0])
            state, reward, done, truncated, info = env.step(action)
            env.render()

            print(f"Reward: {reward}")  # 打印奖励
            print(f"timestep: {env.timestep}")  
            # print(f"Simulation Timestep: {env.get_simulation_timestep()}")  # 打印仿真步进时间
            
            # 导出倒数六个关节的数据
            data = [timestep]
            current_positions = []
            target_positions = []
            kps = []
            kis = []
            kds = []
            torques = []
            for i in range(24, 30):
                current_positions.append(state[i])
                target_positions.append(env.setpoint_curve(env.timestep)[i])
                kps.append(0)  # kp 设置为 0
                kis.append(0)
                kds.append(0)
                torques.append(env.control_output[i])
            data.extend(current_positions)
            data.extend(target_positions)
            data.extend(kps)
            data.extend(kis)
            data.extend(kds)
            data.extend(torques)
            file.write(','.join(map(str, data)) + '\n')
            
            # if done:
            #     state, _ = env.reset()  # 重新初始化环境状态
    env.close()
